# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: daniel

import os
import matplotlib.pyplot as plt
import pandas as pd
from absl import app, flags


FLAGS = flags.FLAGS


def _sumzip(*items):
  return [sum(values) for values in zip(*items)]


class RateLimitHitAnalyzer(object):
  def __init__(self):
    self._rate_limit_hit_summary = []

  def filter_rate_limit_report_records(self, records):
    ret = []
    for record in records:
      if record['type'] != 'OG_LOG':
        continue
      if record['og_log']['type'] != 'RATE_LIMIT_REPORT_MESSAGE':
        continue
      ret.append(record)
    return ret

  def retrieve_rate_limit_hit_record(self,
                                     records,
                                     action,
                                     start_dt,
                                     end_dt):
    tmp = []
    for record in records:
      if record['og_log']['rate_limit_report_message']['last_hit_request_type'] != action:
        continue
      rec = {
          'timestamp': int(record['og_log']['rate_limit_report_message']['last_hit_timestamp']),
          'num_of_hit': 1,
      }
      tmp.append(rec)
    num = len(tmp)
    ret = []
    x_beg = int(pd.to_datetime(start_dt.strftime('%Y%m%dT%H%M%S')).timestamp()) * 1000000000
    x_end = int(pd.to_datetime(end_dt.strftime('%Y%m%dT%H%M%S')).timestamp()) * 1000000000
    x = x_beg
    y = 0
    while x_beg < x_end:
      for each in tmp:
        if each['timestamp'] < x_beg:
          continue
        if each['timestamp'] >= x_beg + 300000000000:
          continue
        y = y + 1
      info = {'timestamp': x, 'num_of_hit': y}
      ret.append(info)
      x_beg = x_beg + 300000000000
      x = x_beg
      y = 0

    return num, sorted(ret, key=lambda x: x['timestamp'])

  def generate_rate_limit_hit_summary_report(self,
                                             trade_date,
                                             output_dir,
                                             local_run,
                                             slurm_run):
    if slurm_run:
      # another process flow for slurm run
      return
    if len(self._rate_limit_hit_summary) == 0:
      return
    records = sorted(self._rate_limit_hit_summary, key=lambda x: x['total_num'], reverse=True)
    color = ['steelblue', 'orange', 'chocolate', 'green']
    label = []
    query_color = []
    query_total = []
    place_color = []
    place_total = []
    amend_color = []
    amend_total = []
    cancel_color = []
    cancel_total = []
    for record in records:
      label.append(record['strategy_name'])
      query_total.append(record['query_num'])
      query_color.append(color[0])
      place_total.append(record['place_num'])
      place_color.append(color[1])
      amend_total.append(record['amend_num'])
      amend_color.append(color[2])
      cancel_total.append(record['cancel_num'])
      cancel_color.append(color[3])

    plt.rcParams.update({'font.size': 7})
    fig = plt.figure(figsize=(30, 16), dpi=200)
    # draw bar chart
    plt.subplot()
    plt.bar(label, query_total, width=0.35, label='Query', color=query_color)
    plt.bar(label, place_total, width=0.35, label='Place', color=place_color,
            bottom=_sumzip(query_total))
    plt.bar(label, amend_total, width=0.35, label='Amend', color=amend_color,
            bottom=_sumzip(query_total, place_total))
    plt.bar(label, cancel_total, width=0.35, label='Cancel', color=cancel_color,
            bottom=_sumzip(query_total, place_total, amend_total))
    plt.ylabel('total # of hits')
    plt.xticks(range(len(label)), label, rotation=270)

    plt.title('rate limit hit summary')
    plt.grid(True)
    plt.legend()
    fig.tight_layout()
    if local_run is True:
      plt.show()
    else:
      output_path = f"{output_dir}/{trade_date}/rate_limit_hit"
      if not os.path.exists(output_path):
        os.makedirs(output_path)
      output_file = f"{trade_date}_rate_limit_hit_summary_report.svg"
      print(f"write to {output_path}/{output_file}")
      fig.savefig(f"{output_path}/{output_file}", dpi=fig.dpi)
    plt.close(fig)

  def generate_rate_limit_hit_report(self,
                                     records,
                                     strategy_name,
                                     machine,
                                     start_dt,
                                     end_dt,
                                     output_dir,
                                     local_run):
    rate_limit_report_records = self.filter_rate_limit_report_records(records)
    # ignore strategy with zero rate limit hit
    if len(rate_limit_report_records) == 0:
      return

    rate_limit_hit_query_num, rate_limit_hit_query = \
        self.retrieve_rate_limit_hit_record(rate_limit_report_records, 'QUERY', start_dt, end_dt)
    rate_limit_hit_place_num, rate_limit_hit_place = \
        self.retrieve_rate_limit_hit_record(rate_limit_report_records, 'PLACE', start_dt, end_dt)
    rate_limit_hit_amend_num, rate_limit_hit_amend = \
        self.retrieve_rate_limit_hit_record(rate_limit_report_records, 'AMEND', start_dt, end_dt)
    rate_limit_hit_cancel_num, rate_limit_hit_cancel = \
        self.retrieve_rate_limit_hit_record(rate_limit_report_records, 'CANCEL', start_dt, end_dt)

    rate_limit_hit_total_num = rate_limit_hit_query_num + \
                               rate_limit_hit_place_num + \
                               rate_limit_hit_amend_num + \
                               rate_limit_hit_cancel_num
    name = strategy_name
    info = {
        'strategy_name': name,
        'total_num': rate_limit_hit_total_num,
        'query_num': rate_limit_hit_query_num,
        'place_num': rate_limit_hit_place_num,
        'amend_num': rate_limit_hit_amend_num,
        'cancel_num': rate_limit_hit_cancel_num,
    }
    self._rate_limit_hit_summary.append(info)

    label = ['Query', 'Place', 'Amend', 'Cancel']
    color = ['steelblue', 'orange', 'chocolate', 'green']
    total = [
        rate_limit_hit_query_num,
        rate_limit_hit_place_num,
        rate_limit_hit_amend_num,
        rate_limit_hit_cancel_num,
    ]
    explode = [0, 0, 0, 0]
    plt.rcParams.update({'font.size': 7})
    fig = plt.figure(figsize=(10, 6), dpi=200)
    # draw pie chart
    plt.subplot(221)
    plt.pie(total, colors=color, explode=explode, labels=label, shadow=True, autopct='%1.1f%%')

    # draw bar chart
    plt.subplot(222)
    plt.bar(label, total, width=0.5, label='', color=color)
    plt.ylabel('# of hits')

    # draw linear chart
    plt.subplot(212)
    if rate_limit_hit_query_num > 0:
      df = pd.DataFrame(rate_limit_hit_query)
      df.timestamp = pd.to_datetime(df.timestamp, unit='ns')
      plt.plot(df.timestamp,
               df.num_of_hit,
               marker='.',
               markersize=0.1,
               label='Query',
               color=color[0])

    if rate_limit_hit_place_num > 0:
      df = pd.DataFrame(rate_limit_hit_place)
      df.timestamp = pd.to_datetime(df.timestamp, unit='ns')
      plt.plot(df.timestamp,
               df.num_of_hit,
               marker='.',
               markersize=0.1,
               label='Place',
               color=color[1])

    if rate_limit_hit_amend_num > 0:
      df = pd.DataFrame(rate_limit_hit_amend)
      df.timestamp = pd.to_datetime(df.timestamp, unit='ns')
      plt.plot(df.timestamp,
               df.num_of_hit,
               marker='.',
               markersize=0.1,
               label='Amend',
               color=color[2])

    if rate_limit_hit_cancel_num > 0:
      df = pd.DataFrame(rate_limit_hit_cancel)
      df.timestamp = pd.to_datetime(df.timestamp, unit='ns')
      plt.plot(df.timestamp,
               df.num_of_hit,
               marker='.',
               markersize=0.1,
               label='Cancel',
               color=color[3])

    plt.title('# of rate limit hit in 5 mins')
    plt.grid(True)
    plt.legend()
    fig.tight_layout(pad=1.0, w_pad=1.0, h_pad=2.0)
    if local_run is True:
      plt.show()
    else:
      trade_date = start_dt.strftime('%Y%m%d')
      start_dt_str = start_dt.strftime('%Y%m%dT%H%M%S')
      end_dt_str = end_dt.strftime('%Y%m%dT%H%M%S')
      output_path = f"{output_dir}/{trade_date}/rate_limit_hit/{strategy_name}/{machine}"
      if not os.path.exists(output_path):
        os.makedirs(output_path)
      output_file = f"{start_dt_str}-{end_dt_str}.svg"
      print(f"write to {output_path}/{output_file}")
      fig.savefig(f"{output_path}/{output_file}", dpi=fig.dpi)
    plt.close(fig)


def main():
  # TODO(daniel):
  print("@daniel: to impliment")


if __name__ == '__main__':
  flags.DEFINE_string('root_dir',
                      '/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log',
                      'root dir specify')
  flags.DEFINE_string('output_dir',
                      '/remote/iosg/strat-1/buckets/analysis.derived.coin/live/plot/og_log_plot',
                      'output_dir.')
  flags.DEFINE_string('tmp_dir',
                      '/tmp',
                      'tmp dir for slurm run')
  flags.DEFINE_string('kafka_config_filename', None, 'kafka config')
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')
  flags.DEFINE_string('machine', None, 'Instance machine name.')
  flags.DEFINE_string('strategy_name', None, 'Strategy name.')
  flags.DEFINE_string('start_time', None, '%Y%m%dT%H%M%S')
  flags.DEFINE_string('end_time', None, '%Y%m%dT%H%M%S')
  flags.DEFINE_string('symbol', None, '')
  flags.DEFINE_string('log_type', None, '')
  flags.DEFINE_bool('local_run', False, "")
  flags.DEFINE_bool('slurm_run', False, "")
  flags.DEFINE_bool('generate_summary_report', False, "generate summary report only")
  app.run(main)
